#include "state.hpp"

using namespace QuantumWalk;

template<typename T> 
state<T>::state(const state<T>& psi){
  if(this != &psi){
    data=psi.data;
    dimV=psi.dimV;
    dimD=psi.dimD;
  }
}

template<typename T> 
state<T>& state<T>::operator=(const state<T>& psi){
  if(this != &psi){
    data=psi.data;
    dimV=psi.dimV;
    dimD=psi.dimD;
  }
  return *(this);
}

template<typename T>
state<T>::state(matrix<T> mat){
  int nrow=mat.size1();
  int ncol=mat.size2();
  dimV=ncol;
  dimD=nrow;
  std::pair<typename std::map<int,vector<T> >::iterator,bool> ret;
  for(int k=0; k < nrow ; k++){
    ret=data.insert(std::pair<int,vector<T> > (k,row(mat,k)));
  
    if(ret.second ==false){
      //Element existed
      std::cout<<"Element "<<k<<" existed:"<<std::endl;
      std::cout<<data[k]<<std::endl;
      std::cout<<"It is now replaced by: "<<std::endl;
      std::cout<<row(mat,k)<<std::endl;
      
      data.erase(ret.first);
      data.insert(std::pair<int,vector<T> > (k,row(mat,k)));
    }
  }
}


template<typename T> 
state<T>& state<T>::operator*=(shift<T>& s_op){
  
  
  
  typename std::map<int,matrix<T> >::iterator op_it;
  typename std::map<int,vector<T> >::iterator state_it;
  
  for(state_it = data.begin();state_it != data.end();state_it++){
    
    op_it=s_op.op_.find(state_it->first);
    //Check
    if(op_it==s_op.op_.end()){
      std::cerr<<"Internal degree "<<state_it->first<<" is not found in shift operator"<<std::endl;
      exit(1);
    }
    
    if(op_it->second.size1() != (unsigned) dimV){
      std::cerr<<"Dimension of shift (permutation) operator and number of sites do not match! Exiting..."
	       <<std::endl;
      exit(1);
    }
    
    //Multiplication
    state_it->second=prod(op_it->second,state_it->second);
  }
  return *this;
}


template<typename T>
state<T>& state<T>::operator*=(const coin<T>& C){
  //Create local copy
  state<T> s0=*this;
  
  if(dimD != C.get_dim()){
    std::cerr<<"Dimension of coin and internal degrees of state do not match! Exiting..."<<std::endl;
    exit(1);
  }
  
  typename std::map<int,vector<T> >::iterator s_it1,s_it2,it_find;
  
  //Iterate over internal degrees of freedom twice
  for(s_it1 = s0.data.begin();s_it1 != s0.data.end();s_it1++){
    //s_it->first is k, where k is in D.
    //Clear the values of the map element
    it_find=data.find(s_it1->first);
    
    //Check that the element is found!
    if(it_find == data.end()){
      std::cerr<<"Could not find element with key: "<<it_find->first<<std::endl;
      exit(1);
    }
    
    //Uses boost::ublas vector clear method.
    it_find->second.clear();
    
    for(s_it2 = s0.data.begin();s_it2 != s0.data.end();s_it2++){
      //s_it2->first is l, where l is in D.
      //Now create the superposition of different spin states.
      //For single site in internal state  k it means \sum_j c_{k,j}\alpha_{j,v}$
      //Note that s_it2->second points to vector \f$ (\alpha_{j,1},\alpha_{j,2},...,\alpha_{j,N}) \f$.
      it_find->second += C(s_it1->first,s_it2->first)*(s_it2->second);
    }
  }
  return *this;
}

template<typename T>
vector<double> state<T>::prob_sites(){
  vector<double> ret(dimV,0);
  typename std::map<int,vector<T> >::iterator it;
  for(int k=0; k < dimV; k++){
    for(it = data.begin();it != data.end(); it++){
      ret(k)+=std::pow(std::abs(it->second(k)),2);
    }
  }
  return ret;
}

template<typename T> 
vector<double> state<T>::prob_internal(){
  vector<double> ret(dimD,0);
  typename std::map<int,vector<T> >::iterator it;
  for(it = data.begin();it != data.end(); it++){
    for(int k=0; k < dimV; k++){
      ret(it->first)+=std::pow(std::abs(it->second(k)),2);
    }
  }
  return ret;
}


template<typename T>
const T& state<T>::operator()(int d,int v) const {
  
  typename std::map<int,vector<T> >::const_iterator it ;
  it=data.find(d);
  
  if(it == data.end()){
    std::cerr<<"Element not found"<<std::endl;
    exit(1);
  }
  return it->second(v);
}

template class state<std::complex<double> >;
//template class state<double>;

//template state<std::complex<double> >o perator*(const shift<std::complex<double> > S,
// 						state<std::complex<double > > psi);

// template state<std::complex<double> > operator*(const coin<std::complex<double> > C,
//  						state<std::complex<double > > psi);

// template state<std::complex<double> > operator*(const shift<double> S,state<double> psi);

// template state<double> operator*(const coin<double> C,state<double> psi);

